﻿#include "kcftracker.hpp"
#include <unistd.h>
#include <ros/ros.h>  
#include <image_transport/image_transport.h> //用来在ROS系统中的话题上发布和订阅图象消息
#include <sensor_msgs/image_encodings.h>
#include <cv_bridge/cv_bridge.h>     //cv_bridge中包含CvBridge库
#include <sensor_msgs/image_encodings.h> //ROS图象类型的编码函数
#include <kcf_tracker/kcfTrackerRoi.h>
#include <kcf_tracker/M_KCF_Select.h>
using namespace std;
using namespace cv;

string uavname = "uav1";


// KCF tracker 初始化参数
bool HOG = true;
bool FIXEDWINDOW = false;
bool MULTISCALE = true;
bool SILENT = true;
bool LAB = false;

Mat videoFrame;
Rect result;

Rect2d trackerRoi_Rect2d;
bool is_tracking = false;
bool is_recTrackerRoi = false;

float response_value =0;
int trackFailCnt=0;


kcf_tracker::kcfTrackerRoi trackerResultRoi; //仅用于记录框的位置
kcf_tracker::M_KCF_Select trackerResultRoi_serial; //仅用于记录框的位置
cv::Mat colorImg ;
bool colorImgRec_Flag = false;
cv::Mat depthImg,depthImg_copy;
bool depthImgRec_Flag = false;

void trackerRoi_cb(const kcf_tracker::kcfTrackerRoi::ConstPtr& msg)
{
	if(msg->state == 0)
		is_tracking = false;
	else
	{
		trackerRoi_Rect2d.x = msg->x;
		trackerRoi_Rect2d.y = msg->y;
		trackerRoi_Rect2d.width = msg->width;
		trackerRoi_Rect2d.height = msg->height;
		is_recTrackerRoi = true;
	}
}
void trackerRoi_serial_cb(const kcf_tracker::M_KCF_Select::ConstPtr& msg)
{
	if(msg->state == 0)
	{
		is_tracking = false;
	}
	else if(msg->state == 1)
	{
		trackerRoi_Rect2d.x = msg->x;
		trackerRoi_Rect2d.y = msg->y;
		trackerRoi_Rect2d.width = msg->width;
		trackerRoi_Rect2d.height = msg->height;
		is_recTrackerRoi = true;
	}
}
void colorImg_cb(const sensor_msgs::ImageConstPtr& msg)
{
    try
    {
        colorImg = cv_bridge::toCvCopy(msg,"bgr8")->image;
        colorImgRec_Flag = true ;
    }
    catch (cv_bridge::Exception& e)
    {
        cout << "colorImg_cb could not convert from " << msg->encoding.c_str() << "to 'brg8'." << endl;
    }   
}

void depthImg_cb(const sensor_msgs::ImageConstPtr& msg)
{
    try
    {
        depthImg = cv_bridge::toCvCopy(msg, sensor_msgs::image_encodings::TYPE_16UC1)->image;
        depthImgRec_Flag = true ;
    }
    catch (cv_bridge::Exception& e)
    {
        cout << "depthImg_cb could not convert from " << msg->encoding.c_str() << "to 'brg8'." << endl;
    }   
}

int main(int argc, char* argv[])
{	
    ros::init(argc, argv,"kcf_tracker_node");
    ros::NodeHandle nh;//创建句柄
	ros::NodeHandle nh_private("~");
    image_transport::ImageTransport it(nh);

	image_transport::Publisher sourceImg_pub = it.advertise("/kcfTracker/image_raw",1);
	image_transport::Subscriber colorImg_sub = it.subscribe("/camera/color/image_raw",1,colorImg_cb);;
	image_transport::Subscriber depthImg_sub = it.subscribe("/camera/aligned_depth_to_color/image_raw", 1, depthImg_cb);
	// ros::Subscriber trackerRoi_sub = nh.subscribe<kcf_tracker::kcfTrackerRoi>("uav1/set/trackRoi", 1, trackerRoi_cb);
	ros::Subscriber trackerRoi_sub = nh.subscribe<kcf_tracker::M_KCF_Select>("/serial_p9/receive/kcf_select", 1, trackerRoi_serial_cb);
	ros::Publisher trackerResultRoi_pb = nh.advertise<kcf_tracker::kcfTrackerRoi>("/kcfTracker/result/trackRoi", 1);
	ros::Publisher trackerResultRoi_serial_pb = nh.advertise<kcf_tracker::M_KCF_Select>("/serial_p9/send/kcf_select", 1);

	double response_threshold_=0;
	int trackFailCnt_threshold_=0;
  	nh_private.param<double>("response_threshold", response_threshold_, 0.3);
  	nh_private.param<int>("trackFailCnt_threshold", trackFailCnt_threshold_, 20);
	nh_private.param<bool>("MULTISCALE", MULTISCALE, true);
	
	cout << endl;
	cout << "response_threshold_     :  " <<  response_threshold_  << endl;
	cout << "trackFailCnt_threshold_ :  " <<  trackFailCnt_threshold_  << endl;
	cout << endl;

	// Create KCFTracker object
	KCFTracker tracker(HOG, FIXEDWINDOW, MULTISCALE, LAB);

	ros::Rate loop_rate(50);
	while (ros::ok())
	{
		if(is_recTrackerRoi)
		{
			is_recTrackerRoi = false;
			tracker.init(trackerRoi_Rect2d, videoFrame);
			rectangle(videoFrame, trackerRoi_Rect2d, Scalar(0, 255, 0), 2, 8);
			is_tracking = true;
		
		}
		if(depthImgRec_Flag)
		{
			depthImgRec_Flag = false;
			depthImg.copyTo(depthImg_copy);
		}
		if(colorImgRec_Flag)
		{
			colorImgRec_Flag = false;
			colorImg.copyTo(videoFrame);
			if (is_tracking == true)
			{
				if (videoFrame.rows == 0 || videoFrame.cols == 0)
				{
					cout << "\033[K" << "videoFrame is error" << endl;
					continue;
				}
				else
				{
					result = tracker.update(videoFrame, response_value);
				
					trackerResultRoi.x = max(result.x,0);
					trackerResultRoi.y = max(result.y,0);
					trackerResultRoi.width = result.x + result.width - max(result.x,0);
					trackerResultRoi.height = result.y + result.height - max(result.y,0);
					trackerResultRoi.state = 1;

					trackerResultRoi_serial.x = trackerResultRoi.x;
					trackerResultRoi_serial.y = trackerResultRoi.y;
					trackerResultRoi_serial.width = trackerResultRoi.width;
					trackerResultRoi_serial.height = trackerResultRoi.height;
					trackerResultRoi_serial.state = 1;

					Point point1 = Point(trackerResultRoi.x + trackerResultRoi.width/4,trackerResultRoi.y + trackerResultRoi.height/4);
					Point point2 = Point(trackerResultRoi.x + trackerResultRoi.width*3/4,trackerResultRoi.y + trackerResultRoi.height*3/4);
					rectangle(videoFrame, point1, point2, Scalar(255, 0, 0), 1, LINE_4);
					if(response_value <= response_threshold_)
					{
						trackFailCnt++;
						cout << "\033[K" << "trackFailCnt:  " << trackFailCnt << "     " << response_value << endl;
						trackerResultRoi.state = 2;
						trackerResultRoi_serial.state = 2;
						
						if(trackFailCnt > trackFailCnt_threshold_)
						{
							trackFailCnt = 0;
							is_tracking = false;
							trackerResultRoi.state = 0;

							//告知地面站
							trackerResultRoi_serial.state = 0;
							trackerResultRoi_serial_pb.publish(trackerResultRoi_serial);

							cout << "\033[K" << "track Fail ！！！" << endl;
						}	
					}
					else
						trackFailCnt = 0;

					switch (trackerResultRoi_serial.state)
					{
					case 1:
						rectangle(videoFrame, Point(result.x, result.y), Point(result.x + result.width, result.y + result.height), Scalar(0, 255, 0), 2, 8);
						break;
					case 2:
						rectangle(videoFrame, Point(result.x, result.y), Point(result.x + result.width, result.y + result.height), Scalar(0, 0, 255), 2, 8);
						break;					
					default:
						break;
					}
					double distance = 0;
					int cnt = 0;
					for (int i = trackerResultRoi.y + trackerResultRoi.height/4 ; i < min(depthImg_copy.rows,trackerResultRoi.y + trackerResultRoi.height*3/4); i++)
					{
						for (int j = trackerResultRoi.x + trackerResultRoi.width/4; j <  min(depthImg_copy.cols, trackerResultRoi.x + trackerResultRoi.width*3/4); j++)
						{
							distance += depthImg_copy.at<u_int16_t>(i,j);
							cnt++ ;
						}
					}
					distance = 0.001 * distance / cnt;
					trackerResultRoi.distance = distance;
					trackerResultRoi.header.stamp = ros::Time::now();
					trackerResultRoi_pb.publish(trackerResultRoi);

					cout << "\033[K" << "\033[32m" << "Center_x	Center_y	distance" << "\033[0m" << endl;
					cout << setiosflags(ios::fixed) << setprecision(5)	<< "\033[K"						
						 << trackerResultRoi.x + trackerResultRoi.width / 2.0  << "\t"
						 << trackerResultRoi.y + trackerResultRoi.height / 2.0  << "\t"																
						 << distance << endl;
					cout << "\033[3A" << endl;
				}		
			}
			// imshow("Image", videoFrame);
			// waitKey(1);
			sensor_msgs::ImagePtr sourceImg = cv_bridge::CvImage(std_msgs::Header(),"bgr8",videoFrame).toImageMsg();
			sourceImg_pub.publish(sourceImg);
		}

        ros::spinOnce();
        loop_rate.sleep(); 
	}
	cout << "\033[2J" << endl;
	return 0;
}




